from Model_Simulation import * 
import numpy as np 
from tensorforce.agents import Agent
from utils_agent import * 
from Parameters_1D import *
totalDepth, axialNodes, totalNodes = spatialVariables_1D_act() 

LAI_factors = np.loadtxt(".../relevant_data/LAI_factors_reduced.txt")

# #Load the agent
# agent_name = 'Scheduler_agent -1-1'
# agent_loaded = Agent.load("./Trained_Agent/",filename=agent_name)
# internals = agent_loaded.initial_internals()

# #Define some simulation metrics and relevant arrays
# episode_len = 30 
# stateArray = np.zeros((episode_len+1, axialNodes+4))
# stateArray_rz = np.zeros(episode_len+1)
# actionsArray = np.zeros(episode_len)

# #Define the bounds of the zone
# lower_zone = 0.216
# upper_zone = 0.280

# def compute_rootZone_moisture(current_state):
#         if current_state[-2] == 0.50:
#             rootZone_vol_moist = 0.10*((1/6)*(current_state[10] + current_state[11] + current_state[12] + current_state[13] + current_state[14] + current_state[15])) +\
#                          0.20*((1/6)*(current_state[15] + current_state[16] + current_state[17] + current_state[18] + current_state[19] + current_state[20])) +\
#                          0.30*((1/6)*(current_state[20] + current_state[21] + current_state[22] + current_state[23] + current_state[24] + current_state[25])) + \
#                          0.40*((1/6)*(current_state[25] + current_state[26] + current_state[27] + current_state[28] + current_state[29] + current_state[30]))
#         else:
#             rootZone_vol_moist = 0.10*((1/6)*(current_state[0] + current_state[1] + current_state[2] + current_state[3] + current_state[4] + current_state[5]))  + \
#                          0.20*((1/6)*(current_state[5] + current_state[6] + current_state[7] + current_state[8] + current_state[9] + current_state[10])) + \
#                          0.30*((1/11)*(current_state[10] + current_state[11] + current_state[12] + current_state[13] + current_state[14] + current_state[15] + current_state[16] + current_state[17] + current_state[18] + current_state[19] + current_state[20])) + \
#                          0.40*((1/11)*(current_state[20] + current_state[21] + current_state[22] + current_state[23] + current_state[24] + current_state[25] + current_state[26] + current_state[27] + current_state[28] + current_state[29] + current_state[30]))

#         return rootZone_vol_moist

# #Set the initial state
# current_rootDepth = 1.00
# rooting_depths = np.zeros(episode_len+1)
# rooting_depths[0:30] = 1.00
# rooting_depths[30:] = 1.00

# lai_factors = LAI_factors[0:episode_len+1]

# state_init = np.zeros(axialNodes+4)
# state_init[0:axialNodes] = 0.22*np.ones(axialNodes)
# state_init[axialNodes]   =  np.random.uniform(1.04, 9.0)
# state_init[axialNodes+1] = np.random.uniform(0.20, 1.25)
# state_init[axialNodes+2] = rooting_depths[0]
# state_init[axialNodes+3] = LAI_factors[0]
# stateArray_rz[0] = compute_rootZone_moisture(state_init)
# stateArray[0,:] = state_init

# for i in range(episode_len):
#     #Generate some noise and add same to the states and weather data
#     # noise_states = 0*np.random.uniform(0.01, 0.05, size=(axialNodes)) # Using a uniform distribution
#     # noise_weather = np.zeros(2)
#     # noise_all = np.concatenate((noise_states,noise_weather))
#     # states = stateArray[i,:] + noise_all
#     states = stateArray[i,:] 

#     for j in range(axialNodes):
#         if states[j] < 0.078:
#             states[j] = 0.078
#             pass
#         if states[j] > 0.43:
#             states[j] = 0.43
#             pass
#     actions, internals = agent_loaded.act(states=states, internals=internals, independent=True)
#     actionsArray[i] = -1*actions["irrigationRate"]*actions["irrigationDecision"]
#     current_state=stateArray[i,:]
#     next_state = SimulateModel(current_state[0:axialNodes], actions["irrigationRate"]*actions["irrigationDecision"],
#                                current_state[axialNodes], current_state[axialNodes+1], current_state[-2], current_state[-1])
    
#     next
    
#     #next_state[-1] = current_rootDepth
#     next_state[-2] = rooting_depths[i+1]
#     next_state[-1] = lai_factors[i+1]
#     stateArray_rz[i+1] = compute_rootZone_moisture(next_state)
#     stateArray[i+1,:] = next_state
#     pass 

# #Generate and save a plot
# filename = 'Trained_Agent_Res_3'
# GenerateTrajectories(episode_len, actionsArray,stateArray_rz,lower_zone, upper_zone,filename)

#Generate a plot of the reward trajectory
filename_reward = 'Reward_Trajectory_5_mz2'
reward_trajectory = np.loadtxt("reward_trajectory_5.txt")
PlotRewardTrajectory(rewardTrajectory=reward_trajectory, filename=filename_reward)
